from typing import Union, Tuple, Optional

from paddle.fluid.dygraph import Sequential
from paddle.fluid.dygraph.nn import Conv2D, BatchNorm, Pool2D
from paddle.fluid.layers import reshape

__all__ = ['BNAct', 'Pool2d', 'GlobalAvgPool', 'set_fp16', "Conv2d", "flatten"]

FP16 = False


def flatten(x):
    c = x.shape[1] * x.shape[2] * x.shape[3]
    return reshape(x, [-1, c])


def set_fp16(mode: bool):
    global FP16
    FP16 = mode


def BNAct(channels, act):
    return BatchNorm(num_channels=channels, act=act)


def Conv2d(in_channels: int,
           out_channels: int,
           kernel_size: Union[int, Tuple[int]],
           stride: Union[int, Tuple[int]] = 1,
           padding: Union[str, int, Tuple[int]] = 'same',
           groups: int = 1,
           dilation: int = 1,
           bias: Optional[bool] = None,
           bn: bool = False,
           act: Optional[str] = None):

    if isinstance(kernel_size, int):
        kernel_size = (kernel_size, kernel_size)

    if padding == 'same':
        kh, kw = kernel_size
        ph = (kh + (kh - 1) * (dilation - 1) - 1) // 2
        pw = (kw + (kw - 1) * (dilation - 1) - 1) // 2
        padding = (ph, pw)

    use_cudnn = FP16 or (out_channels != groups)

    if bn or act is not None:
        if bn:
            conv = Conv2D(
                num_channels=in_channels,
                num_filters=out_channels,
                filter_size=kernel_size,
                stride=stride,
                padding=padding,
                groups=groups,
                use_cudnn=use_cudnn,
            )
            return Sequential(
                conv,
                BNAct(out_channels, act),
            )
        else:
            conv = Conv2D(
                num_channels=in_channels,
                num_filters=out_channels,
                filter_size=kernel_size,
                stride=stride,
                padding=padding,
                groups=groups,
                use_cudnn=use_cudnn,
                act=act,
            )
            return conv
    else:
        return Conv2D(
            num_channels=in_channels,
            num_filters=out_channels,
            filter_size=kernel_size,
            stride=stride,
            padding=padding,
            groups=groups,
            use_cudnn=use_cudnn,
        )

def Pool2d(kernel_size, stride, padding='same', type='avg', ceil_mode=True):
    if isinstance(kernel_size, int):
        kernel_size = (kernel_size, kernel_size)

    if padding == 'same':
        kh, kw = kernel_size
        ph = (kh - 1) // 2
        pw = (kw - 1) // 2
        padding = (ph, pw)

    return Pool2D(
        pool_size=kernel_size,
        pool_stride=stride,
        pool_padding=padding,
        pool_type=type,
        ceil_mode=ceil_mode,
    )


def GlobalAvgPool():
    use_cudnn = FP16
    return Pool2D(pool_type='avg', global_pooling=True, use_cudnn=use_cudnn)
